#!/usr/bin/env python
import sys
import os
import shutil
import cPickle
import re
import time
import glob
from subprocess import *
import zipfile
import itertools
import tempfile
import random

class Trash:
 
 def __init__(self):
  HOME = os.environ['HOME']
  self.trashDirectory = '%s/.TrashPy'%HOME
  c = itertools.count()
  time.sleep(random.random()*2)
  while glob.glob('%s/lock*'%self.trashDirectory) != []:
   sys.stdout.write('The database is locked, waiting the database... %s'%c.next())
   sys.stdout.flush()
   sys.stdout.write("\r")
   time.sleep(random.random()*2)
  self.lock = tempfile.mkdtemp(prefix='lock', dir=self.trashDirectory)
  self.files = sys.argv[2:]
  for file in self.files:
   if not re.findall('\d+\.\d+-\d+.zip', file):
    if not (os.path.isfile(file) or os.path.isdir(file)):
     print "No such file or directory: %s"%file
     os.rmdir(self.lock)
     sys.exit()

  try:
   self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'r')
   self.filesId = cPickle.load(self.trashdbf)
   self.trashdbf.close()
  except IOError:
   self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
   self.filesId = {}
   cPickle.dump(self.filesId, self.trashdbf)

 def help(self):
  print "put\t: put the given file or directory to Trash\nlist\t: list all the file or directory into Trash\nclist\t: Contextual list: list all items in the trash corresponding to the current working directory\nrestore\t: restore the given file or directory\nempty\t: empty the trash\ncempty\t: Contextual empty: only remove items in the trash coming from the current working directory\nrm\t: remove the given file from the trash"

 def put(self):
  append = False
  timestamp=str(time.time())
  key = '%s.zip'%(timestamp)
  cwd = os.getcwd()
  dirPrec = ''
  keysDirs = {}
  toRemove = {}
  c = itertools.count()
  for path in self.files:
   path = re.sub('/$', '', path)
   file = os.path.basename(path)
   dir = path.replace(file, '')
   if dir == '':
    dir = cwd
   os.chdir(dir)
   dir = os.getcwd()
   os.chdir(cwd)
   if dir != dirPrec:
    timestamp='%s-%s'%(time.time(), c.next())
    key = '%s.zip'%(timestamp)
    toRemove.update({key:[file]})
    keysDirs.update({key:dir})
   else:
    toRemove[key].append(file)
   dirPrec = dir
  for key in toRemove.keys():
   dir = keysDirs[key]
   files = toRemove[key]
   os.chdir(dir)
   try:
    os.mkdir('.trash')
   except OSError:
    pass
   for file in files:
    os.system('cp -rvL %s .trash && rm -rv %s'%(file,file))
   os.chdir(cwd)
  for key in toRemove.keys():
   dir = keysDirs[key]
   files = toRemove[key]
   os.chdir('%s/.trash'%dir)
#   z = zipfile.ZipFile(key, 'w', allowZip64=True)
#   z.write(files[0])
   self.zip(files[0], key)
   os.system('rm -rvf %s'%(files[0]))
   for file in files[1:]:
#    z.write(file)
    self.zip(file, key)
    os.system('rm -rvf %s'%(file))
#   z.close()
   self.filesId.update({key: dir})
   self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
   cPickle.dump(self.filesId, self.trashdbf)
   self.trashdbf.close()
   os.chdir(cwd)
#  self.gzipAndMv(keysDirs)

 def zip(self, folder_path, output_path):
  """Zip the contents of an entire folder (with that folder included
  in the archive). Empty subfolders will be included in the archive
  as well.
  """
  if os.path.isdir(folder_path):
   parent_folder = os.path.dirname(folder_path)
   # Retrieve the paths of the folder contents.
   contents = os.walk(folder_path)
   try:
    zip_file = zipfile.ZipFile(output_path, 'a', zipfile.ZIP_DEFLATED, allowZip64=True)
    for root, folders, files in contents:
     # Include all subfolders, including empty ones.
     for folder_name in folders:
      absolute_path = os.path.join(root, folder_name)
      relative_path = absolute_path.replace(parent_folder + '\\',
                                            '')
      print "Adding '%s' to archive." % absolute_path
      zip_file.write(absolute_path, relative_path)
     for file_name in files:
      absolute_path = os.path.join(root, file_name)
      relative_path = absolute_path.replace(parent_folder + '\\',
                                            '')
      print "Adding '%s' to archive." % absolute_path
      zip_file.write(absolute_path, relative_path)
    print "'%s' created successfully." % output_path
   except IOError, message:
    print message
    sys.exit(1)
   except OSError, message:
    print message
    sys.exit(1)
   except zipfile.BadZipfile, message:
    print message
    sys.exit(1)
   finally:
    zip_file.close()
  elif os.path.isfile(folder_path):
   zip_file = zipfile.ZipFile(output_path, 'a', zipfile.ZIP_DEFLATED, allowZip64=True)
   zip_file.write(folder_path)
   zip_file.close()

  
 def list(self):
  cwd = os.getcwd()
  for key in sorted(self.filesId.keys()):
   dir = self.filesId[key]
   try:
    os.chdir('%s/.trash'%dir)
    z = zipfile.ZipFile(key, 'r', allowZip64=True)
    namelist = z.namelist()
    print '\033[1;32m%s\t:\t%s\n%s files in archive\033[1;m' % (key, dir, len(namelist))
    for f in namelist[:11]:
     print f
    os.chdir(cwd)
   except OSError:
    print "No such file or directory: %s/.trash"%dir
    print "Removing %s from database"%key
    self.filesId.pop(key)
    self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
    cPickle.dump(self.filesId, self.trashdbf)
    self.trashdbf.close()
   except IOError:
    print "No such file or directory: %s/.trash/%s"%(dir, key)
    print "Removing %s from database"%key
    self.filesId.pop(key)
    self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
    cPickle.dump(self.filesId, self.trashdbf)
    self.trashdbf.close()

 def cList(self):
  """
  Contextual list: display only item which correspond to the current working directory
  """
  cwd = os.getcwd()
  for key in sorted(self.filesId.keys()):
   if re.findall('%s$'%cwd.strip(), '%s'%self.filesId[key].strip()):
    dir = self.filesId[key]
    try:
     os.chdir('%s/.trash'%dir)
     z = zipfile.ZipFile(key, 'r', allowZip64=True)
     namelist = z.namelist()
     print '\033[1;32m%s\t:\t%s\n%s files in archive\033[1;m' % (key, dir, len(namelist))
     for f in namelist[:11]:
      print f
     os.chdir(cwd)
    except OSError:
     print "No such file or directory: %s/.trash"%dir
     print "Removing %s from database"%key
     self.filesId.pop(key)
     self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
     cPickle.dump(self.filesId, self.trashdbf)
     self.trashdbf.close()
    except IOError:
     print "No such file or directory: %s/.trash/%s"%(dir, key)
     print "Removing %s from database"%key
     self.filesId.pop(key)
     self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
     cPickle.dump(self.filesId, self.trashdbf)
     self.trashdbf.close()


 def cEmpty(self):
  """
  Remove files from the trash only for items which correspond to the current working directory
  """
  cwd = os.getcwd()
  for key in sorted(self.filesId.keys()):
   if re.findall('%s$'%cwd.strip(), '%s'%self.filesId[key].strip()):
    dir = self.filesId[key]
    path = '%s/.trash/%s'%(dir, key)
    print 'rm %s'%path
    try:
     shutil.rmtree(path)
    except OSError:
     try:
      os.remove(path)
     except OSError:
      pass
    self.filesId.pop(key)
    try:
     os.rmdir('%s/.trash'%dir)
    except OSError:
     pass
  self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
  cPickle.dump(self.filesId, self.trashdbf)
  self.trashdbf.close()
 
 def restore(self):
  for file in self.files:
   #os.renames('%s/%s'%(self.trashDirectory, file), '%s/%s'%(self.filesId[file], file))
   dir = self.filesId[file]
   path = '%s/.trash/%s'%(dir, file)
   os.system('mv -v %s %s/.'%(path, self.filesId[file]))
   cwd = os.getcwd()
   os.chdir(self.filesId[file])
   try:
    os.rmdir('.trash')
   except OSError:
    print '.trash directory is not empty'
   z = zipfile.ZipFile(file, 'r', allowZip64=True)
   z.extractall()
   os.system('rm -vf %s'%(file))
   os.chdir(cwd)
   self.filesId.pop(file)
  self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
  cPickle.dump(self.filesId, self.trashdbf)
  self.trashdbf.close()
 
 def empty(self):
  for file in self.filesId.keys():
   dir = self.filesId[file]
   path = '%s/.trash/%s'%(dir, file)
   print 'rm %s'%path
   try:
    shutil.rmtree(path)
   except OSError:
    try:
     os.remove(path)
    except OSError:
     pass
   try:
    os.rmdir('%s/.trash'%dir)
   except:
    pass
   self.filesId.pop(file)
  self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
  cPickle.dump(self.filesId, self.trashdbf)
  self.trashdbf.close()

 def rm(self):
  for files in self.files:
   dir = self.filesId[files]
   for path in glob.glob('%s/.trash/%s'%(dir, files)):
    file = path.split('/')[-1]
    os.system('rm -v %s'%path)
   self.filesId.pop(files)
   try:
    os.rmdir('%s/.trash'%dir)
   except OSError:
    pass
  self.trashdbf = open('%s/.trashdb'%(self.trashDirectory), 'w')
  cPickle.dump(self.filesId, self.trashdbf)
  self.trashdbf.close()

def main():
 trash = Trash()
 if re.findall('help', sys.argv[1]) or sys.argv[1]=='-h':
  trash.help()
 if sys.argv[1] == 'put':
  trash.put()
 if sys.argv[1] == 'list':
  trash.list()
 if sys.argv[1] == 'clist':
  trash.cList()
 if sys.argv[1] == 'restore':
  trash.restore()
 if sys.argv[1] == 'empty':
  trash.empty()
 if sys.argv[1] == 'cempty':
  trash.cEmpty()
 if sys.argv[1] == 'rm':
  trash.rm()
 os.rmdir(trash.lock)

if __name__ == "__main__":
 main()
